package com.milvus.imagesearch;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonPrimitive;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.collection.response.ListCollectionsResp;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.response.InsertResp;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.springframework.stereotype.Service;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
@Service
public class MilvusManager {


    private  MilvusClientV2 client=new MilvusClientV2(
            ConnectConfig.builder()
                    .uri("http://116.62.142.81:19530")
                    .token("root:Milvus")
                    .build()
    );



    public void createCollection( String collectionName) {


//        DropCollectionReq collectionNameReq = DropCollectionReq.builder()
//                .collectionName(collectionName)
//                .build();
//        client.dropCollection(collectionNameReq);
        ListCollectionsResp listCollectionsResp = client.listCollections();
        if (CollectionUtil.isNotEmpty(listCollectionsResp.getCollectionNames()) &&listCollectionsResp.getCollectionNames().contains(collectionName)){
            return;
        }
        //创建模式
        CreateCollectionReq.CollectionSchema schema = client.createSchema();

        //        // 定义集合和字段
        // 3.2 添加字段
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(true)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("embedding")
                .dataType(DataType.FloatVector)
                .dimension(1000)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("name")
                .dataType(DataType.VarChar)
                .maxLength(512)
                .build());



        // 创建索引
        IndexParam indexParamForIdField = IndexParam.builder()
                .fieldName("id")
                .indexType(IndexParam.IndexType.STL_SORT)
                .build();

        IndexParam indexParamForVectorField = IndexParam.builder()
                .fieldName("embedding")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .metricType(IndexParam.MetricType.COSINE)
                .build();

        List<IndexParam> indexParams = new ArrayList<>();
        indexParams.add(indexParamForIdField);
        indexParams.add(indexParamForVectorField);


        // 3.4 创建具有模式和索引参数的集合
        CreateCollectionReq customizedSetupReq1 = CreateCollectionReq.builder()
                .collectionName(collectionName)
                .collectionSchema(schema)
                .indexParams(indexParams)
                .build();

        client.createCollection(customizedSetupReq1);

        // 3.5 获取集合的加载状态
        GetLoadStateReq customSetupLoadStateReq1 = GetLoadStateReq.builder()
                .collectionName(collectionName)
                .build();

        Boolean loaded = client.getLoadState(customSetupLoadStateReq1);
        System.out.println(loaded);

    }

    public void insertData(String collectionName, INDArray features, String fileName) {
        // 插入数据
        float[] floatArray = features.toFloatVector();

        // 手动将 float[] 转换为 List<Float>
        List<String> floatList = new ArrayList<>();
        for (float f : floatArray) {
            floatList.add(new BigDecimal(f).toString()); // 将每个 float 添加到列表中
        }


        Gson gson = new Gson();
        List<JsonObject> data=new ArrayList<>();
        HashMap<String, List<String>> embedding = new HashMap<String, List<String>>() {{
            put("embedding", floatList);
        }};
        String string = JSONObject.toJSONString(embedding);
        JsonObject jsonObject = gson.fromJson(string, JsonObject.class);

        jsonObject.add("name",new JsonPrimitive(fileName));
        data.add(jsonObject);

        InsertReq insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(data)
                .build();
        InsertResp insertResp = client.insert(insertReq);

    }

}
